#!/usr/bin/env python
# -*- coding: utf-8 -*-

from fastapi import FastAPI, WebSocket, Request, Cookie, Query, status, Depends
from fastapi.responses import HTMLResponse
from fastapi.templating import Jinja2Templates
from typing import Optional




app = FastAPI(debug=True)
templates = Jinja2Templates(directory="./templates")



@app.get("/")
async def home(request: Request):
    return templates.TemplateResponse(
        # "webchat.html",  # 初始
        # "webchat2.html",  # 一对一聊天
        "webchat3.html",  # 一对一聊天（接受/发送json格式消息）
        {
            "request": request
        }
    )




async def get_cookie_or_token(
        websocket: WebSocket,
        session: Optional[str] = Cookie(None),
        token: Optional[str] = Query(None),
):
    if session is None and token is None:
        await websocket.close(code=status.WS_1008_POLICY_VIOLATION)
    return session or token


from connection_tool import ConnectionManager
from starlette.websockets import WebSocketDisconnect
ws_manager = ConnectionManager()


@app.websocket("/item/ws")
async def websocket_endpoint(
        websocket: WebSocket,
        q: Optional[int] = None,
        cookie_or_token: str = Depends(get_cookie_or_token),
):
    # await websocket.accept()
    await ws_manager.connect(cookie_or_token, websocket)

    try:
        while True:
            data = await websocket.receive_text()
            if data == 'bye':
                await websocket.send_text("请注意：窗口即将关闭！")
                await websocket.send_text("请注意：窗口已关闭！")
                await websocket.close()
                break
            else:
                # await websocket.send_text(
                #     f"Session cookie or query token value is: {cookie_or_token}"
                # )
                # if q is not None:
                #     await websocket.send_text(f"parameter q is: {q}")
                await websocket.send_text(f"Message is: {data}")
    except WebSocketDisconnect as e:
        await ws_manager.disconnect(cookie_or_token, websocket)




# WebSockets（六）聊天室完善:增加广播功能、发送所有人消息、发送个人消息

@app.websocket("/item/ws/{user}")
async def websocket_chat_user(
        websocket: WebSocket,
        user: str,
        cookie_or_token: str = Depends(get_cookie_or_token)
):
    await ws_manager.connect(user, websocket)
    await ws_manager.broadcast(user + "进入聊天室")
    try:
        while True:
            data = await websocket.receive_text()
            await ws_manager.send_personal_message(f"你说：{data}", websocket)
            # await ws_manager.broadcast(f"用户{user}说：{data}")
            await ws_manager.broadcast_other(f"用户{user}说：{data}", user)
    except WebSocketDisconnect as e:
        await ws_manager.disconnect(user, websocket)
        await ws_manager.broadcast(f"用户{user}离开")



# WebSockets(七)实现一对一聊天
import json

@app.websocket("/ws/{user}")
async def websocket_one_to_one(
        websocket: WebSocket,
        user: str,
        cookie_or_token: str = Depends(get_cookie_or_token)
):
    """一对一聊天"""
    await ws_manager.connect(user, websocket)
    try:
        while True:
            data = await websocket.receive_text()   # 前端发送的数据是：ws.send(input.value+"?"+username.value)
            await ws_manager.send_other_message(message=json.loads(data.split("?")[0]), user=data.split("?")[1])
    except WebSocketDisconnect as e:
        await ws_manager.disconnect(user, websocket)


# WebSockets(八)接受/发送json格式消息
@app.websocket("/ws/json/{user}")
async def websocket_json_data(
        websocket: WebSocket,
        user: str,
        cookie_or_token: str = Depends(get_cookie_or_token)
):
    """发送接收json数据（前面一对一接口其实就是发送的json数据）"""
    await ws_manager.connect(user, websocket)
    try:
        while True:
            data = await websocket.receive_json()
            send_user = data["username"]
            if send_user:
                # await ws_manager.send_other_message(data, send_user)
                data["username"] = '你'
                await ws_manager.send_personal_message_json(data, websocket)
                data["username"] = user
                await ws_manager.send_other_message(data, send_user)
            else:
                data["username"] = user+'对全体'
                await ws_manager.broadcast_json(data)
    except WebSocketDisconnect as e:
        await ws_manager.disconnect(user, websocket)




# WebSockets（三）登录后才可以聊天

from sqlalchemy.orm import Session
from models.models import *
from models.schemas import *
from fastapi import HTTPException

from models.database import SessionLocal


def create_db():
    """每个请求处理完毕后关闭当前连接，不同的请求使用不同的链接"""
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()

def get_user_by_email(db: Session, email: str):
    user = db.query(User).filter(User.email == email).first()
    if not user:
        raise HTTPException(status_code=404, detail="this email not exists")
    return user


@app.get("/login")
async def login(request: Request):
    return templates.TemplateResponse(
        "login.html",
        {
            "request": request
        }
    )


@app.post("/token")
def generate_token(
        user: UserModel,
        db: Session = Depends(create_db)
):
    db_user = get_user_by_email(db, user.email)
    client_hash_password = user.password + "_hashed"
    if client_hash_password == db_user.hashed_password:
        return {"token": db_user.email, "msg": "success"}
    return {"token": None, "msg": "failed"}


if __name__ == '__main__':
    import uvicorn
    uvicorn.run("main:app", host="127.0.0.1", port=8000, reload=True)